# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

from collections import namedtuple
import datetime
import functools
import logging
import os
import signal

from absl import app, flags
from tornado.ioloop import IOLoop

import coin.flow.subscriber as flow_subscriber
from coin.base.param_util import to_list
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_order_checker_pb2 import AccountCheckerConfig
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.hk_hard_arb.pricer import PositionLeanPricer
from coin.strategy.hk_hard_arb.index_pricer import IndexPricer
from coin.strategy.hk_hard_arb.util import (load_clean_json, parse_subreq_config2)
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.passive.util import (convert_executor_config_from_dict)
from coin.support.accounting.checker.account_checker import (AccountChecker,
                                                             AccountCheckFailedError)
from coin.support.accounting.checker.config_util import convert_checker_config_from_dict
from coin.proto.coin_order_gateway_pb2 import OrderGatewayConfig
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.strategy.mm.subscription import FeedSubscriptionRequest

FLAGS = flags.FLAGS

flags.DEFINE_string('og_configs', '', '')

flags.DEFINE_string('config_filename', '', '')

flags.DEFINE_integer('exit_after_min', 30, '')

flags.DEFINE_boolean('dry_run', False, '')

MMProduct = namedtuple('MMProduct', ['shared', 'trade', 'ref'])


def _gen_mm_products(products_config, ref_product_type, trade_product_type):
  mm_products = []
  shared_symbol_set = set()
  for shared_symbol, prod_config in products_config.items():
    assert 'trade_symbol' in prod_config and 'ref_symbol' in prod_config
    assert shared_symbol not in mm_products
    shared_symbol_set.add(shared_symbol)
    trade_prod = trade_product_type.FromStrProduct(prod_config['trade_symbol'])
    ref_symbol = to_list(prod_config['ref_symbol'])
    assert isinstance(ref_product_type, (list, tuple))
    assert len(ref_symbol) == len(ref_product_type)
    ref_prod = [
        prod_type.FromStrProduct(symbol) for symbol, prod_type in zip(ref_symbol, ref_product_type)
    ]
    mm_products.append(MMProduct(shared=shared_symbol, trade=trade_prod, ref=ref_prod))
  return mm_products


def _gen_for_same_quotes(config):
  amt_det_notional = config['price_det_notional']
  book_askt_func_0 = (
      lambda book: book.get_notional_asks_by_amt(amt_det_notional).entries[-1].price)
  book_askt_func_1 = book_askt_func_0
  book_bidt_func_0 = (
      lambda book: book.get_notional_bids_by_amt(amt_det_notional).entries[-1].price)
  book_bidt_func_1 = book_bidt_func_0
  return book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1


def _gen_for_different_quotes(config):
  assert 'price_det_notional' not in config
  assert 'ref_qty_det_notional' in config or 'ref_price_det_notional' in config
  assert 'trade_qty_det_notional' in config or 'trade_price_det_notional' in config

  if 'ref_qty_det_notional' in config:
    qty_det_notional = to_list(config['ref_qty_det_notional'])
    book_askt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_asks_by_qty(notional).entries[-1].price, det)
        for det in qty_det_notional
    ]
    book_bidt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_bids_by_qty(notional).entries[-1].price, det)
        for det in qty_det_notional
    ]
  elif 'ref_price_det_notional' in config:
    amt_det_notional = to_list(config['ref_price_det_notional'])
    book_askt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_asks_by_amt(notional).entries[-1].price, det)
        for det in amt_det_notional
    ]
    book_bidt_func_0 = [
        functools.partial(
            lambda notional, book: book.get_notional_bids_by_amt(notional).entries[-1].price, det)
        for det in amt_det_notional
    ]
  else:
    raise ValueError('invalid ref')

  # TODO(chensili): trade_qty_det_notional
  if 'trade_qty_det_notional' in config:
    qty_det_notional = config['trade_qty_det_notional']
    book_askt_func_1 = (
        lambda book: book.get_notional_asks_by_qty(qty_det_notional).entries[-1].price)
    book_bidt_func_1 = (
        lambda book: book.get_notional_bids_by_qty(qty_det_notional).entries[-1].price)
  elif 'trade_price_det_notional' in config:
    amt_det_notional = config['trade_price_det_notional']
    book_askt_func_1 = (
        lambda book: book.get_notional_asks_by_amt(amt_det_notional).entries[-1].price)
    book_bidt_func_1 = (
        lambda book: book.get_notional_bids_by_amt(amt_det_notional).entries[-1].price)
  else:
    raise ValueError('invalid ref')
  return book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1


def _gen_true_book_funcs(config):
  if 'price_det_notional' in config:
    return _gen_for_same_quotes(config)
  else:
    return _gen_for_different_quotes(config)


class SingleProductSubStrategy(object):
  def __init__(self, *, ref_product, trade_product, ccy_0_1, parent, config, logger=None):
    self._trade_product = trade_product
    self._ref_product = to_list(ref_product)
    self._products = self._ref_product + to_list(self._trade_product)

    self._logger = logger

    self._edge = None
    self._spread = None
    self._pull_edge = None

    if 'edge_bp' in config:
      self._edge = config['edge_bp'] / 10000.
    if 'spread_bp' in config:
      self._spread = config['spread_bp'] / 10000.
    if 'pull_edge_bp' in config:
      self._pull_edge = config['pull_edge_bp'] / 10000.
    if 'pull_spread_bp' in config:
      self._pull_edge = config['pull_spread_bp'] / 10000. / 2

    assert 1 == sum([
        self._edge is not None,
        self._spread is not None,
    ])

    self._bookmap = BookMap(self._products, logger=self._logger)

    book_askt_func_0, book_askt_func_1, book_bidt_func_0, book_bidt_func_1 = \
        _gen_true_book_funcs(config)

    self._pricer = IndexPricer(basis_ma_window=(int(config['basis_ma_window_min'] * 60e9)),
                               tick=get_holder_from_product(trade_product).product_info.gen_tick(),
                               ref_book_askt_funcs=book_askt_func_0,
                               ref_book_bidt_funcs=book_bidt_func_0,
                               trade_book_askt_func=book_askt_func_1,
                               trade_book_bidt_func=book_bidt_func_1,
                               price_multiplier_1=ccy_0_1)
    self._pos_lean_pricer = PositionLeanPricer(min_pos=config['min_pos'],
                                               max_pos=config['max_pos'],
                                               lot_size=config['lot_size'],
                                               lean_edge_bp=config.get('lean_edge_bp', 0))

    # composite instances
    self._parent = parent
    self._executor = None
    self._config = config

  def on_og_reset(self):
    # Intialize passive executor
    self._logger.info('Intializing passive executor...')
    execution_params = convert_executor_config_from_dict(self._trade_product.symbol, self._config)

    self._executor = construct_mm_executor(self._parent._trade_subreq,
                                           execution_params,
                                           self._parent.og,
                                           self._logger)

  def on_book(self, exchange_id, book):
    product = self._products[exchange_id]

    self._bookmap.set_book(product, book)

    # Check book freshness and sanity
    if not self._bookmap.is_book_ready() or not self._bookmap.is_book_sane():
      return

    ref_book = [self._bookmap.get_book(product) for product in self._ref_product]
    trade_book = self._bookmap.get_book(self._trade_product)
    self._pricer.update(ref_book, trade_book)

    if not self._pricer.ready or self._executor is None or FLAGS.dry_run:
      return

    if self._spread is not None:
      askp, bidp = self._pricer.calculate_price_by_spread(self._spread)
    elif self._edge is not None:
      askp, bidp = self._pricer.calculate_price(self._edge, self._edge)
    else:
      raise ValueError('not supported mode')

    pos = self._executor.query_unbiased_position()
    askp, bidp = self._pos_lean_pricer.update(pos=pos, buy_price=bidp, sell_price=askp)
    askp, bidp = self._pricer.adjust_price(askp_1=askp, bidp_1=bidp)

    sell_price_pull = None
    buy_price_pull = None
    if askp is not None and self._pull_edge is not None:
      sell_price_pull = askp * (1 - self._pull_edge)
    if bidp is not None and self._pull_edge is not None:
      buy_price_pull = bidp * (1 + self._pull_edge)

    self._executor.manage_mm_orders(sell_price=askp,
                                    buy_price=bidp,
                                    sell_price_pull=sell_price_pull,
                                    buy_price_pull=buy_price_pull)


class PassUnhedgeStrategy(object):
  def __init__(self, strategy_params):
    self._logger = logging.getLogger(__name__)

    self._config = strategy_params['config']

    self._trade_subreq, self._trade_feed_subreq, self._ref_feed_subreqs = \
        parse_subreq_config2(self._config['common'])

    self._exchange = self._trade_subreq.exchange
    self._og_config = strategy_params['og_config']

    self._balance_checker = None
    account_checker_config = self._get_account_checker_config()
    if account_checker_config is not None:
      self._balance_checker = AccountChecker(account_checker_config)

    # TODO(someone): Need some unified metaclass
    ref_book_builder = [get_flow_book_builder_cls(req) for req in self._ref_feed_subreqs]
    trade_book_builder = get_flow_book_builder_cls(self._trade_feed_subreq)
    self._book_builder_clses = ref_book_builder + to_list(trade_book_builder)

    self._live_env_clses = (
        None,
        get_live_env_cls(self._trade_subreq),
    )
    self._live_env = None

    self.og = None

    self._single_product_sub_strategies = []

    ref_product_type = [book_builder.ProductType for book_builder in self._book_builder_clses[:-1]]
    trade_product_type = self._book_builder_clses[-1].ProductType
    self._mm_products = _gen_mm_products(self._config['products'],
                                         ref_product_type,
                                         trade_product_type)

    ref_exchange_feed_checker = self._config['common'].get('ref_exchange_feed_checker', False)
    trade_exchange_feed_checker = self._config['common'].get('trade_exchange_feed_checker', False)
    feed_checker_products = []
    if ref_exchange_feed_checker:
      feed_checker_products += self._get_trade_products()
    if trade_exchange_feed_checker:
      feed_checker_products += self._get_ref_products()
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)

    self._feed = MultiFeedSubsystem(feed_checker_config=feed_checker_config)

    self._balance_printer = None

    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None

    self._ccy_0_1 = self._config['common'].get('ccy_0_1', None) or 1

  def _get_trade_currencies(self):
    trade_currencies = \
      [mm_product.trade.base for mm_product in self._mm_products]
    return trade_currencies

  def _get_trade_products(self):
    trade_products = [mm_product.trade for mm_product in self._mm_products]
    return trade_products

  def _get_ref_products(self):
    ref_products = []
    for mm_product in self._mm_products:
      for product in to_list(mm_product.ref):
        ref_products.append(product)
    return ref_products

  def _get_account_checker_config(self):
    if self._config.get('account_checker') is None:
      return None

    config_keys = {
        'reserve',
        'lot_size',
        'stack_size',
        'ratio_min_balance',
        'ratio_max_balance',
        'ratio_min_hold_balance',
    }

    balance_check_config_list = []
    for currency, config in self._config['account_checker'].items():
      config = {**config}
      if 'lot_size' not in config:
        for product in self._config['products']:
          base = SharedProduct.FromStrNativeProduct(product).base.currency
          if base == currency:
            config['lot_size'] = self._config['products'][product]['lot_size']
            break

      assert config_keys == set(config.keys()), config
      reserve = config['reserve']
      config['currency'] = currency
      config['min_balance'] = reserve * config['ratio_min_balance']
      config['max_balance'] = reserve * config['ratio_max_balance']
      config['min_available_balance'] = 0.0
      config['max_available_balance'] = config['lot_size'] * config['stack_size']
      config['min_hold_balance'] = reserve * config['ratio_min_hold_balance']
      config['max_hold_balance'] = reserve * config['ratio_max_balance']
      currency_checker_config = convert_checker_config_from_dict(config)
      balance_check_config_list.append(currency_checker_config)

    return AccountCheckerConfig(each_balance_check_config=balance_check_config_list)

  def subscribe_feeds(self):
    common_config = self._config['common']

    for product_id, mm_product in enumerate(self._mm_products):
      shared_symbol = mm_product.shared
      config = self._config['products'][shared_symbol]
      config = {**common_config, **config}
      self._single_product_sub_strategies.append(
          SingleProductSubStrategy(ref_product=mm_product.ref,
                                   trade_product=mm_product.trade,
                                   ccy_0_1=self._ccy_0_1,
                                   parent=self,
                                   config=config,
                                   logger=self._logger))
      for exchange_id, product in enumerate(mm_product.ref):
        self._feed.subscribe_book(
            FeedSubscriptionRequest.create([product], self._ref_feed_subreqs[exchange_id]),
            functools.partial(self._on_book, product_id, exchange_id))
      trade_product_exchange_id = len(self._ref_feed_subreqs)
      self._feed.subscribe_book(
          FeedSubscriptionRequest.create([mm_product.trade], self._trade_feed_subreq),
          functools.partial(self._on_book, product_id, trade_product_exchange_id))

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    og_params = self._config.get('og_params', {})
    og_params = convert_og_config_from_dict(self._trade_subreq, og_params)
    self._live_env = create_live_env_instance(self._live_env_clses[1],
                                              self._get_trade_currencies(),
                                              self._get_trade_products(),
                                              self._og_config,
                                              og_params=og_params,
                                              ioloop=ioloop,
                                              queue=queue,
                                              flow_sub=flow_sub)
    self._live_env.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def register_order_gateway(self, name, order_gateway):
    self._logger.info('OG initialized: %s', str(name))
    self.og = order_gateway
    # cancel working orders first after OGs are ready.
    if not self._order_canceled_after_og_ready:
      self.og.cancel_all()
      self._logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True

    for sub_strategy in self._single_product_sub_strategies:
      sub_strategy.on_og_reset()

    self._exit_executor = SimpleCancelAllExecutor()

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    self._live_env.start()
    self.print_balance(periodic=True)

  def _on_book(self, product_id, exchange_id, book):
    if not self.og:
      self._logger.info('OG not ready')
      return
    if self._exit_attempt > 0:
      return

    self._single_product_sub_strategies[product_id].on_book(exchange_id, book)

  def print_balance(self, periodic=False):
    if self._balance_printer is not None:
      self._balance_printer.print_balance()

    if self._balance_checker is not None and self.og is not None and self.og.is_ready():
      balance_proto = self.og.og_info.balance_as_proto()
      try:
        self._balance_checker.check_balances_by_range(balance_proto)
      except AccountCheckFailedError:
        self._logger.exception('Balance is wrong, kill strategy.\n%s', balance_proto)
        os.kill(os.getpid(), signal.SIGINT)  # commit suicide

    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.print_balance, True)

  def signal_handler(self, sig, frame):
    self._logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None:
      ogs = to_list(self.og)
      self._logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)


def get_param_set(og_configs, config_filename):
  og_config = og_configs.split(',')[0]
  config = load_clean_json(config_filename)
  return {
      'og_config': og_config,
      'config': config,
  }


def main(argv):
  strategy_params = get_param_set(og_configs=FLAGS.og_configs,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('unhedged_mm')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()

  strategy = PassUnhedgeStrategy(strategy_params)
  strategy.subscribe_feeds()
  strategy.subscribe_live_envs(ioloop, queue, flow_sub)  # What the..
  strategy.start(ioloop, queue, flow_sub)

  exit_after = datetime.timedelta(minutes=FLAGS.exit_after_min)
  ioloop.add_timeout(exit_after, strategy.on_exit)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  ioloop.start()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
